package com.yvon.boot.mybatis.utils;

import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import lombok.AllArgsConstructor;

import java.io.Serializable;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;

/**
 * @author : Yvon
 * @since : 2021-09-26
 */
@AllArgsConstructor
public class TreeSearch<T> {

    /**
     * 业务搜素节点的方法
     */
    private Function<Serializable, T> selectById;

    /**
     * 业务节点设置子节点
     */
    private BiConsumer<T, List<T>> setChildren;

    /**
     * 业务节点获取子节点
     */
    private Function<T, List<T>> getChildren;

    /**
     * 获取节点ID的方法
     */
    private Function<T, Serializable> getId;

    /**
     * 获取父节点ID的方法
     */
    private Function<T, Serializable> getParentId;

    /**
     * 判断根节点的方法
     */
    private Predicate<T> isRoot;


    public List<T> search(List<T> flatNodes) {
        if (CollectionUtils.isEmpty(flatNodes)) {
            return flatNodes;
        }

        final List<T> roots = new LinkedList<>();

        final Map<Serializable, T> nodes = new HashMap<>();
        flatNodes.forEach(t -> nodes.put(getId.apply(t), t));
        List<T> pList = flatNodes;
        Set<Serializable> rids = new HashSet<>();
        Serializable tid;
        while (!pList.isEmpty()) {
            flatNodes = pList;
            pList = new LinkedList<>();
            for (T t : flatNodes) {
                if (isRoot.test(t) && !rids.contains((tid = getId.apply(t)))) {
                    roots.add(t);
                    rids.add(tid);
                }
            }
        }
        return null;
    }
}
